#include <iostream>
#include <AscendAI/AscendAI.h>
#include <fstream>
#include <vector>
#include <cmath>

void load_input_data(const std::string& filename, std::vector<float>& data, size_t& rows, size_t& cols) {
    // 读取输入数据文件 (.npy)
    std::ifstream infile(filename, std::ios::binary);
    // 假设文件格式是平坦的二进制数据
    infile.read(reinterpret_cast<char*>(&rows), sizeof(rows));
    infile.read(reinterpret_cast<char*>(&cols), sizeof(cols));
    data.resize(rows * cols);
    infile.read(reinterpret_cast<char*>(data.data()), data.size() * sizeof(float));
}

void save_output_data(const std::vector<float>& data, size_t rows, size_t cols, const std::string& filename) {
    // 保存计算结果为 .npy 文件
    std::ofstream outfile(filename, std::ios::binary);
    outfile.write(reinterpret_cast<const char*>(&rows), sizeof(rows));
    outfile.write(reinterpret_cast<const char*>(&cols), sizeof(cols));
    outfile.write(reinterpret_cast<const char*>(data.data()), data.size() * sizeof(float));
}

int main() {
    // 初始化 Ascend SDK
    AscendAI::init();

    // 加载输入数据
    std::vector<float> data;
    size_t rows = 0, cols = 0;
    load_input_data("input_data.npy", data, rows, cols);

    // 将数据转化为张量
    AscendAI::Tensor input_tensor(rows, cols, data.data());

    // 示例操作：矩阵乘法（这里用假数据表示）
    AscendAI::Tensor result_tensor = AscendAI::matrix_multiply(input_tensor, input_tensor); // 假设进行矩阵乘法

    // 获取计算结果并保存
    std::vector<float> result_data = result_tensor.toVector();
    save_output_data(result_data, rows, cols, "output_result.npy");

    std::cout << "Computation completed and result saved as 'output_result.npy'." << std::endl;

    return 0;
}
